# -*- coding: utf-8 -*-
"""
中性策略框架 | 邢不行 | 2024分享会
author: 邢不行
微信: xbx6660
"""
import time
import traceback

import pandas as pd
from utils.commons import retry_wrapper, truncate
from api.market import fetch_binance_ticker_data


# =====获取持仓
# 获取币安账户的实际持仓
def get_position_df(exchange):
    """
    获取币安账户的实际持仓

    :param exchange:        交易所对象，用于获取数据
    :return:

              当前持仓量   均价  持仓盈亏
    symbol
    RUNEUSDT       -82.0  1.208 -0.328000
    FTMUSDT        523.0  0.189  1.208156

    """
    # 获取原始数据
    position_df = retry_wrapper(exchange.fapiPrivateV2_get_positionrisk, params={'timestamp': ''}, func_name='获取账户持仓风险')
    position_df = pd.DataFrame(position_df, dtype='float')    # 将原始数据转化为dataframe

    # 整理数据
    columns = {'positionAmt': '当前持仓量', 'entryPrice': '均价', 'unRealizedProfit': '持仓盈亏', 'markPrice': '当前标记价格'}
    position_df.rename(columns=columns, inplace=True)
    position_df = position_df[position_df['当前持仓量'] != 0]  # 只保留有仓位的币种
    position_df.set_index('symbol', inplace=True)  # 将symbol设置为index

    # 保留指定字段
    position_df = position_df[['当前持仓量', '均价', '持仓盈亏', '当前标记价格']]

    return position_df


def get_equity(exchange):
    """
    获取账户净值

    :param exchange:    交易所对象

    :return:
        swap_equity=1000  (表示账户里资金总价值为 1000U )

    """
    # 获取U本位合约账户净值(不包含未实现盈亏)
    equity = retry_wrapper(exchange.fapiPrivateV2_get_account, params={'timestamp': ''}, func_name='获取U本位合约账户净值')  # 获取账户净值
    equity = pd.DataFrame(equity['assets'])
    usdt_balance = float(equity[equity['asset'] == 'USDT']['walletBalance'])  # 获取usdt资产
    equity = usdt_balance

    return equity


def get_spot_position(exchange):
    """
    获取账户净值

    :param exchange:    交易所对象

    :return:
        swap_equity=1000  (表示账户里资金总价值为 1000U )

    """
    # 获取U本位合约账户净值(不包含未实现盈亏)
    position_df = retry_wrapper(exchange.private_get_account, params={'timestamp': ''}, func_name='获取现货账户净值')  # 获取账户净值
    position_df = pd.DataFrame(position_df['balances'], dtype=float)
    position_df['free'] += position_df['locked']
    position_df = position_df[position_df['free'] != 0]

    position_df.rename(columns={'asset': 'symbol', 'free': '当前持仓量'}, inplace=True)

    # 保留指定字段
    position_df = position_df[['symbol', '当前持仓量']]
    position_df['仓位价值'] = None  # 设置默认值

    return position_df


def get_spot_position_and_equity(exchange):
    # =获取现货持仓
    spot_position = get_spot_position(exchange)
    # 获取现货最新价格
    spot_price = fetch_binance_ticker_data(exchange, symbol_type='spot')
    if spot_position.empty:
        spot_usdt = 0
        spot_equity = 0
        dust_spot = pd.DataFrame()
    else:
        if 'USDT' in spot_position['symbol'].to_list():
            spot_usdt = spot_position.loc[spot_position['symbol'] == 'USDT', '当前持仓量'].iloc[0]  # 获取当前账号U的数量
        else:
            spot_usdt = 0
        spot_position.loc[spot_position['symbol'] != 'USDT', 'symbol'] = spot_position['symbol'] + 'USDT'  # 追加USDT后缀，方便计算usdt价值
        has_price_spot = list(set(spot_position['symbol'].to_list()) & set(spot_price.index))  # 保留含有USDT报价的现货
        spot_position = spot_position[spot_position['symbol'].isin(has_price_spot)]  # 过滤掉没有报价的现货
        spot_position.set_index('symbol', inplace=True)
        spot_position['当前价格'] = spot_price
        spot_position['仓位价值'] = spot_position['当前持仓量'] * spot_position['当前价格']
        # BNB在现货持仓需要剔除
        if 'BNBUSDT' in spot_position.index:
            # BNB会去购买合约，不会存在现货仓位
            spot_position.drop('BNBUSDT', inplace=True)
        dust_spot = spot_position[spot_position['仓位价值'] < 5]  # 无法下单的碎单
        spot_position = spot_position[spot_position['仓位价值'] > 5]  # 过滤掉一些无法下的碎单，不计入持仓价值中

        spot_equity = spot_position['仓位价值'].sum() + spot_usdt  # 计算当前现货账户总价值，按U计算

    return spot_equity, spot_position, spot_usdt, dust_spot


def bn_transfer_u_to_spot(exchange, amount, asset='USDT'):
    """
    BN U本位合约账户 转U 到 现货账户
    :param exchange:    交易所对象
    :param amount:      划转资金
    :param asset:       资产类型，默认：USDT
    """
    print('bn_transfer_u_to_spot：', amount, asset)
    if amount <= 0:
        print(f'当前划转金额【{amount}】小于0，本次不进行划转。(不是错误不要担心)')
        return
    params = {
        'type': 'UMFUTURE_MAIN',
        'asset': asset,
        'amount': str(amount),
        'timestamp': '',
    }
    transfer_info = retry_wrapper(exchange.sapi_post_asset_transfer, params=params, func_name='U本位合约账户转U去现货账户')
    print(f'BN 资金转移 {asset} U本位合约账户 => 现货账户 接口返回: ', transfer_info)
    print(f'BN 资金转移 {asset} U本位合约账户 => 现货账户 资金量: ', amount)


def bn_transfer_spot_to_u(exchange, amount, asset='USDT'):
    """
    BN 现货账户 转U 到 U本位合约账户
    :param exchange:    交易所对象
    :param amount:      划转资金
    :param asset:       资产类型，默认：USDT
    :return:
    """
    print('bn_transfer_spot_to_u', amount, asset)
    if amount <= 0:
        print(f'当前划转金额【{amount}】小于0，本次不进行划转。(不是错误不要担心)')
        return

    if exchange.is_pure_long and asset == 'USDT':
        print('纯多策略不划转钱到合约')
        return

    params = {
        'type': 'MAIN_UMFUTURE',
        'asset': asset,
        'amount': str(amount),
        'timestamp': '',
    }
    transfer_info = retry_wrapper(exchange.sapi_post_asset_transfer, params=params, func_name='BN现货账户转U去U本位合约账户')
    print(f'BN 资金转移 {asset} 现货账户 => U本位合约账户 接口返回: ', transfer_info)
    print(f'BN 资金转移 {asset} 现货账户 => U本位合约账户 资金量: ', amount)


def replenish_bnb(exchange, buy_bnb_value, is_use_spot):
    """
    补充BNB，用于抵扣手续费

    :param exchange:        交易所对象
    :param buy_bnb_value:   补充价值多少U的bnb
    :param is_use_spot:     是否使用现货模式
    """
    # ===获取当前持仓BNB的价值
    # 获取BNB价格
    ticker = retry_wrapper(exchange.public_get_ticker_price, params={'symbol': 'BNBUSDT'}, func_name='获取BNB现货价格')
    price = float(ticker['price'])  # bnb当前现货价格

    # 获取合约BNB数量
    swap_balance = retry_wrapper(exchange.fapiPrivateV2_get_account, params={'timestamp': ''}, func_name='获取U本位合约账户净值')  # 获取账户净值
    swap_balance = pd.DataFrame(swap_balance['assets'])
    swap_bnb = float(swap_balance[swap_balance['asset'] == 'BNB']['walletBalance'])  # 获取BNB持仓数量
    swap_bnb_value = swap_bnb * price  # 计算当前BNB价值

    # 获取现货BNB数量
    spot_balance = retry_wrapper(exchange.private_get_account, params={'timestamp': ''}, func_name='获取现货账户净值')  # 获取账户净值
    spot_balance = pd.DataFrame(spot_balance['balances'], dtype=float)
    spot_bnb = float(spot_balance[spot_balance['asset'] == 'BNB']['free'])  # 获取BNB持仓数量
    spot_bnb_value = spot_bnb * price  # 计算当前BNB价值
    print(f'当前BNB剩余数量：现货 {spot_bnb_value:.4f} U，合约 {swap_bnb_value:.4f} U')

    # ===构建判断条件和需要购买BNB的价值
    if is_use_spot:  # 现货模式，现货和合约账户都需要进行判断是否购买bnb
        condition = swap_bnb_value + spot_bnb_value < buy_bnb_value  # 判断两边BNB总和是否小于购买量
        all_buy_bnb_value = buy_bnb_value - (swap_bnb_value + spot_bnb_value)  # 现货模式，计算购买合约和现货的bnb
    else:  # 合约模式，只需要判断合约账户是否需要购买bnb
        condition = swap_bnb_value < buy_bnb_value  # 判断合约的BNB是否小于购买量
        all_buy_bnb_value = buy_bnb_value - swap_bnb_value  # 合约模式，计算购买合约的bnb

    # ===判断账户BNB是否 小于 buy_bnb_value
    if condition:
        # =判断购买BNB数量是否满足最小下单量
        if all_buy_bnb_value < 10:
            print('补充现货和合约BNB总价值不足10U，跳过补充BNB操作')
            all_buy_bnb_value = 0
        else:
            print(f'补充现货和合约BNB总价值 {all_buy_bnb_value:.4f} U，准备开始补充BNB······')
            try:
                # 划转资金去现货账户
                bn_transfer_u_to_spot(exchange, all_buy_bnb_value)
            except BaseException as e:
                print(traceback.format_exc())
                print('划转U到现货账户报错，准备尝试消耗现货预留手续费资金···')

            # 市价单购买BNB
            ticker = exchange.public_get_ticker_price(params={'symbol': 'BNBUSDT'})
            price = float(ticker['price'])
            quantity = all_buy_bnb_value / price * 0.98  # 相当于2%滑点
            quantity = float(f'{quantity:.3f}')  # 根据价格精度调整，这里直接写死了
            params = {'symbol': 'BNBUSDT', 'side': 'BUY', 'type': 'MARKET', 'quantity': str(quantity), 'timestamp': ''}
            print('购买BNB下单参数: ', params)
            try:
                res = retry_wrapper(exchange.privatePostOrder, params=params, func_name='privatePostOrder')
                print('购买BNB下单结果: ', res)
            except BaseException as e:
                print(traceback.format_exc())
                print('购买BNB出错，跳过补充BNB操作')
                return 0

            # 休息一下，让市场消化一下单子
            time.sleep(5)
    else:
        all_buy_bnb_value = 0

    # 平衡BNB
    rebalance_bnb(exchange, is_use_spot)

    return all_buy_bnb_value


def rebalance_bnb(exchange, is_use_spot):
    """
    根据下单模式，来平衡现货账户和合约账户之间的bnb数量

    :param exchange:        交易所对象
    :param is_use_spot:     是否使用现货模式
    """
    # 获取合约BNB数量
    swap_balance = retry_wrapper(exchange.fapiPrivateV2_get_account, params={'timestamp': ''}, func_name='获取U本位合约账户净值')  # 获取账户净值
    swap_balance = pd.DataFrame(swap_balance['assets'])
    swap_bnb = float(swap_balance[swap_balance['asset'] == 'BNB']['walletBalance'])  # 获取BNB持仓数量

    # 获取现货BNB数量
    spot_balance = retry_wrapper(exchange.private_get_account, params={'timestamp': ''}, func_name='获取现货账户净值')  # 获取账户净值
    spot_balance = pd.DataFrame(spot_balance['balances'], dtype=float)
    spot_bnb = float(spot_balance[spot_balance['asset'] == 'BNB']['free'])  # 获取BNB持仓数量

    # 根据模式，来判断BNB划转
    if is_use_spot:  # 现货模式，设置现货账户与合约账户两边BNB平衡
        bnb_mean = (swap_bnb + spot_bnb) / 2  # 构建两边的平衡点
        transfer_bnb = bnb_mean - swap_bnb  # 需要给合约多少bnb
        transfer_bnb = truncate(transfer_bnb, 6)  # 向下取整，保留6位有效数字，划转金额可以比最小下单精度小，这里多延几位
        if transfer_bnb > 0:  # 若大于0，表示需要向合约划转
            # 划转BNB到合约账户
            bn_transfer_spot_to_u(exchange, transfer_bnb, asset='BNB')
        else:  # 若小于0，表示需要向现货划转
            # 划转BNB到现货账户
            bn_transfer_u_to_spot(exchange, abs(transfer_bnb), asset='BNB')
    else:  # 合约模式，直接将现货账户的BNB全部转入合约账户中
        spot_bnb = truncate(spot_bnb, 6)  # 向下取整，保留6位有效数字，划转金额可以比最小下单精度小，这里多延几位
        if spot_bnb > 0:
            bn_transfer_spot_to_u(exchange, spot_bnb, asset='BNB')
